package org.dxy.trigger;

import javax.swing.Timer;
import javax.swing.*;
import java.awt.*;
import java.awt.event.*;
import java.util.*;
import java.util.List;

import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.factory.Nd4jBackend;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

/**
 * @ description: Snake Game with Deep Q-Learning, fixed hasSafePath for out-of-bounds, added GPU backend verification
 * @ author: dxy
 * @ createTime: 2025/7/18 23:17
 */
public class SnakeGameWithDQN extends JPanel implements ActionListener {
    private final int WIDTH = 600;
    private final int HEIGHT = 600;
    private final int UNIT_SIZE = 20;
    private final int GAME_UNITS = (WIDTH * HEIGHT) / (UNIT_SIZE * UNIT_SIZE);
    private final int DELAY = 100;
    private final ArrayList<Integer> snakeX = new ArrayList<>();
    private final ArrayList<Integer> snakeY = new ArrayList<>();
    private int foodX;
    private int foodY;
    private int score;
    private char direction = 'R';
    private boolean running = false;
    private boolean foodEaten = false;
    private Timer timer;
    private Random random;
    // DQN components
    private MultiLayerNetwork model;
    private Deque<Transition> replayMemory = new LinkedList<>();
    private final int REPLAY_MEMORY_SIZE = 1000;
    private final double GAMMA = 0.99;
    private double epsilon = 1.0;
    private final double EPSILON_MIN = 0.1;
    private final double EPSILON_DECAY = 0.995;
    private final int BATCH_SIZE = 32;
    private int gameCount = 0;
    private double totalReward = 0.0;
    private final Deque<Double> recentRewards = new LinkedList<>();
    private final int REWARD_WINDOW = 10;
    private final double DISTANCE_REWARD_SCALE = 5.0;
    private Deque<Character> recentDirections = new LinkedList<>();
    private int consecutiveFoodMoves = 0;

    public SnakeGameWithDQN() {
        setPreferredSize(new Dimension(WIDTH, HEIGHT));
        setBackground(Color.black);
        setFocusable(true);
        addKeyListener(new MyKeyAdapter());
        // Verify ND4J backend
        Nd4jBackend backend = Nd4j.getBackend();
        System.out.println("ND4J Backend: " + backend.getClass().getName());
        initializeDQN();
        startGame();
    }

    private void initializeDQN() {
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(123)
                .updater(new Adam(0.001))
                .list()
                .layer(new DenseLayer.Builder().nIn(11).nOut(100).activation(Activation.RELU).build())
                .layer(new DenseLayer.Builder().nIn(100).nOut(100).activation(Activation.RELU).build())
                .layer(new OutputLayer.Builder(LossFunctions.LossFunction.MSE)
                        .nIn(100).nOut(4).activation(Activation.IDENTITY).build())
                .build();
        model = new MultiLayerNetwork(conf);
        model.init();
    }

    public void startGame() {
        snakeX.clear();
        snakeY.clear();
        snakeX.add(UNIT_SIZE * 5);
        snakeY.add(UNIT_SIZE * 5);
        score = 0;
        direction = 'R';
        running = true;
        foodEaten = false;
        recentDirections.clear();
        consecutiveFoodMoves = 0;
        random = new Random();
        spawnFood();
        if (timer != null) {
            timer.stop();
        }
        timer = new Timer(DELAY, this);
        timer.start();
        gameCount++;
        recentRewards.add(totalReward);
        if (recentRewards.size() > REWARD_WINDOW) {
            recentRewards.removeFirst();
        }
        totalReward = 0.0;
        System.out.println("Starting game #" + gameCount + ", Epsilon: " + String.format("%.3f", epsilon));
    }

    public void spawnFood() {
        boolean validPosition;
        do {
            validPosition = true;
            foodX = random.nextInt(WIDTH / UNIT_SIZE) * UNIT_SIZE;
            foodY = random.nextInt(HEIGHT / UNIT_SIZE) * UNIT_SIZE;
            if (foodX < 0 || foodX >= WIDTH || foodY < 0 || foodY >= HEIGHT) {
                validPosition = false;
                continue;
            }
            for (int i = 0; i < snakeX.size(); i++) {
                if (foodX == snakeX.get(i) && foodY == snakeY.get(i)) {
                    validPosition = false;
                    break;
                }
            }
        } while (!validPosition);
        System.out.println("Food spawned at (" + foodX + ", " + foodY + ")");
    }

    @Override
    public void paintComponent(Graphics g) {
        super.paintComponent(g);
        draw(g);
    }

    public void draw(Graphics g) {
        // Draw food
        g.setColor(Color.red);
        g.fillOval(foodX, foodY, UNIT_SIZE, UNIT_SIZE);

        // Draw snake
        for (int i = 0; i < snakeX.size(); i++) {
            if (i == 0) {
                g.setColor(Color.green);
            } else {
                g.setColor(new Color(0, 100, 0));
            }
            g.fillRect(snakeX.get(i), snakeY.get(i), UNIT_SIZE, UNIT_SIZE);
        }

        // Draw score and training metrics
        g.setColor(Color.white);
        g.setFont(new Font("Arial", Font.PLAIN, 20));
        g.drawString("Score: " + score + " | Game: " + gameCount, 10, 20);
        g.drawString("Epsilon: " + String.format("%.3f", epsilon), 10, 50);
        double avgReward = recentRewards.isEmpty() ? 0.0 : recentRewards.stream().mapToDouble(Double::doubleValue).average().getAsDouble();
        g.drawString("Avg Reward: " + String.format("%.2f", avgReward), 10, 80);
        g.drawString("Free Space: " + getFreeSpaceCount(), 10, 110);
        g.drawString("Safe Path: " + (hasSafePath() ? "Yes" : "No"), 10, 140);
    }

    public void move() {
        if (foodEaten) {
            snakeX.add(snakeX.get(snakeX.size() - 1));
            snakeY.add(snakeY.get(snakeY.size() - 1));
            foodEaten = false;
        }

        for (int i = snakeX.size() - 1; i > 0; i--) {
            snakeX.set(i, snakeX.get(i - 1));
            snakeY.set(i, snakeY.get(i - 1));
        }

        switch (direction) {
            case 'U':
                snakeY.set(0, snakeY.get(0) - UNIT_SIZE);
                break;
            case 'D':
                snakeY.set(0, snakeY.get(0) + UNIT_SIZE);
                break;
            case 'L':
                snakeX.set(0, snakeX.get(0) - UNIT_SIZE);
                break;
            case 'R':
                snakeX.set(0, snakeX.get(0) + UNIT_SIZE);
                break;
        }
    }

    public double[] getState() {
        double[] state = new double[11];
        int headX = snakeX.get(0);
        int headY = snakeY.get(0);

        state[0] = isDanger(headX, headY - UNIT_SIZE) ? 1 : 0; // Up
        state[1] = isDanger(headX - UNIT_SIZE, headY) ? 1 : 0; // Left
        state[2] = isDanger(headX + UNIT_SIZE, headY) ? 1 : 0; // Right

        state[3] = direction == 'U' ? 1 : 0;
        state[4] = direction == 'D' ? 1 : 0;
        state[5] = direction == 'L' ? 1 : 0;
        state[6] = direction == 'R' ? 1 : 0;

        state[7] = (foodX < headX && foodY < headY) ? 1 : 0; // Left-up
        state[8] = (foodX < headX && foodY > headY) ? 1 : 0; // Left-down
        state[9] = (foodX > headX && foodY < headY) ? 1 : 0; // Right-up
        state[10] = (foodX > headX && foodY > headY) ? 1 : 0; // Right-down

        return state;
    }

    private boolean isDanger(int x, int y) {
        if (x < 0 || x >= WIDTH || y < 0 || y >= HEIGHT) {
            return true;
        }
        for (int i = 1; i < snakeX.size(); i++) {
            if (x == snakeX.get(i) && y == snakeY.get(i)) {
                return true;
            }
        }
        return false;
    }

    public void checkFood() {
        if (snakeX.get(0) == foodX && snakeY.get(0) == foodY) {
            score++;
            foodEaten = true;
            System.out.println("Food eaten! Score: " + score + ", Snake length: " + (snakeX.size() + 1));
            spawnFood();
            consecutiveFoodMoves = 0; // Reset after eating
        }
    }

    public double checkCollision() {
        int headX = snakeX.get(0);
        int headY = snakeY.get(0);
        if (headX < 0 || headX >= WIDTH || headY < 0 || headY >= HEIGHT) {
            running = false;
            System.out.println("Game Over: Wall collision, Penalty: -200.0");
            startGame();
            return -200.0;
        }

        for (int i = 1; i < snakeX.size(); i++) {
            if (headX == snakeX.get(i) && headY == snakeY.get(i)) {
                running = false;
                System.out.println("Game Over: Self collision (ate itself), Penalty: -200.0");
                startGame();
                return -200.0;
            }
        }

        if (!hasSafePath()) {
            running = false;
            System.out.println("Game Over: Trapped in dead-end, Penalty: -150.0");
            startGame();
            return -150.0;
        }

        return -0.05;
    }

    private double getDistanceToFood(int headX, int headY) {
        double distance = Math.abs(headX - foodX) + Math.abs(headY - foodY);
        return Math.max(distance, 1.0); // Avoid division by zero
    }

    private boolean hasSafePath() {
        int headX = snakeX.get(0);
        int headY = snakeY.get(0);
        // Validate head and food coordinates
        if (headX < 0 || headX >= WIDTH || headY < 0 || headY >= HEIGHT ||
                foodX < 0 || foodX >= WIDTH || foodY < 0 || foodY >= HEIGHT) {
            System.out.println("Invalid coordinates: Head (" + headX + ", " + headY + "), Food (" + foodX + ", " + foodY + ")");
            return false;
        }

        boolean[][] visited = new boolean[WIDTH / UNIT_SIZE][HEIGHT / UNIT_SIZE];
        Queue<int[]> queue = new LinkedList<>();
        int gridX = headX / UNIT_SIZE;
        int gridY = headY / UNIT_SIZE;
        queue.add(new int[]{gridX, gridY});
        visited[gridX][gridY] = true;

        int[][] directions = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}}; // Left, Right, Up, Down
        Set<String> snakeBody = new HashSet<>();
        for (int i = 1; i < snakeX.size(); i++) {
            int bodyX = snakeX.get(i) / UNIT_SIZE;
            int bodyY = snakeY.get(i) / UNIT_SIZE;
            snakeBody.add(bodyX + "," + bodyY);
        }

        while (!queue.isEmpty()) {
            int[] pos = queue.poll();
            int x = pos[0];
            int y = pos[1];
            if (x * UNIT_SIZE == foodX && y * UNIT_SIZE == foodY) {
                System.out.println("Safe path found to food at (" + foodX + ", " + foodY + ")");
                return true;
            }

            for (int[] dir : directions) {
                int nx = x + dir[0];
                int ny = y + dir[1];
                if (nx >= 0 && nx < WIDTH / UNIT_SIZE && ny >= 0 && ny < HEIGHT / UNIT_SIZE
                        && !visited[nx][ny] && !snakeBody.contains(nx + "," + ny)) {
                    queue.add(new int[]{nx, ny});
                    visited[nx][ny] = true;
                }
            }
        }
        System.out.println("No safe path to food at (" + foodX + ", " + foodY + ")");
        return false;
    }

    private int getFreeSpaceCount() {
        int headX = snakeX.get(0);
        int headY = snakeY.get(0);
        int count = 0;
        int[][] directions = {{0, -UNIT_SIZE}, {0, UNIT_SIZE}, {-UNIT_SIZE, 0}, {UNIT_SIZE, 0}};
        for (int[] dir : directions) {
            int nx = headX + dir[0];
            int ny = headY + dir[1];
            if (!isDanger(nx, ny)) {
                count++;
            }
        }
        return count;
    }

    private boolean isPredictiveCloser(int action) {
        int headX = snakeX.get(0);
        int headY = snakeY.get(0);
        double oldDistance = getDistanceToFood(headX, headY);
        char tempDirection = direction;
        switch (action) {
            case 0:
                if (direction != 'D') tempDirection = 'U';
                break;
            case 1:
                if (direction != 'U') tempDirection = 'D';
                break;
            case 2:
                if (direction != 'R') tempDirection = 'L';
                break;
            case 3:
                if (direction != 'L') tempDirection = 'R';
                break;
        }
        int newHeadX = headX;
        int newHeadY = headY;
        switch (tempDirection) {
            case 'U': newHeadY -= UNIT_SIZE; break;
            case 'D': newHeadY += UNIT_SIZE; break;
            case 'L': newHeadX -= UNIT_SIZE; break;
            case 'R': newHeadX += UNIT_SIZE; break;
        }
        if (newHeadX < 0 || newHeadX >= WIDTH || newHeadY < 0 || newHeadY >= HEIGHT) {
            return false;
        }
        double newDistance = getDistanceToFood(newHeadX, newHeadY);
        return newDistance < oldDistance;
    }

    private boolean isBacktracking(char newDirection) {
        if (recentDirections.size() < 2) return false;
        List<Character> lastTwo = new ArrayList<>(recentDirections);
        char last = lastTwo.get(lastTwo.size() - 1);
        if ((newDirection == 'U' && last == 'D') || (newDirection == 'D' && last == 'U') ||
                (newDirection == 'L' && last == 'R') || (newDirection == 'R' && last == 'L')) {
            return true;
        }
        return false;
    }

    private int chooseAction(double[] state) {
        if (random.nextDouble() < epsilon) {
            int action = random.nextInt(4);
            System.out.println("Random action: " + action);
            return action;
        }
        INDArray input = Nd4j.create(state).reshape(1, 11);
        INDArray output = model.output(input);
        int action = output.argMax(1).getInt(0);
        System.out.println("Model action: " + action);
        return action;
    }

    private void trainDQN(double[] state, int action, double reward, double[] nextState, boolean done) {
        totalReward += reward;
        replayMemory.add(new Transition(state, action, reward, nextState, done));
        if (replayMemory.size() > REPLAY_MEMORY_SIZE) {
            replayMemory.removeFirst();
        }

        if (replayMemory.size() >= BATCH_SIZE) {
            ArrayList<Transition> batch = new ArrayList<>();
            for (int i = 0; i < BATCH_SIZE; i++) {
                int idx = random.nextInt(replayMemory.size());
                batch.add(new ArrayList<>(replayMemory).get(idx));
            }

            double[][] states = new double[BATCH_SIZE][11];
            double[][] targets = new double[BATCH_SIZE][4];
            int i = 0;
            for (Transition t : batch) {
                states[i] = t.state;
                INDArray target = model.output(Nd4j.create(t.state).reshape(1, 11));
                if (t.done) {
                    target.putScalar(t.action, t.reward);
                } else {
                    INDArray nextQ = model.output(Nd4j.create(t.nextState).reshape(1, 11));
                    double maxNextQ = nextQ.maxNumber().doubleValue();
                    target.putScalar(t.action, t.reward + GAMMA * maxNextQ);
                }
                targets[i] = target.toDoubleVector();
                i++;
            }

            model.fit(Nd4j.create(states), Nd4j.create(targets));
        }

        if (epsilon > EPSILON_MIN) {
            epsilon *= EPSILON_DECAY;
        }
    }

    @Override
    public void actionPerformed(ActionEvent e) {
        if (running) {
            double[] state = getState();
            int action = chooseAction(state);
            int oldHeadX = snakeX.get(0);
            int oldHeadY = snakeY.get(0);
            double oldDistance = getDistanceToFood(oldHeadX, oldHeadY);

            char oldDirection = direction;
            switch (action) {
                case 0:
                    if (direction != 'D') direction = 'U';
                    break;
                case 1:
                    if (direction != 'U') direction = 'D';
                    break;
                case 2:
                    if (direction != 'R') direction = 'L';
                    break;
                case 3:
                    if (direction != 'L') direction = 'R';
                    break;
            }
            if (direction != oldDirection) {
                System.out.println("Action chosen: " + action + " -> Direction: " + direction);
                recentDirections.add(direction);
                if (recentDirections.size() > 2) recentDirections.removeFirst();
            }

            double reward = -0.05; // Default move penalty
            if (isBacktracking(direction)) {
                reward -= 10.0;
                System.out.println("Penalty: -10.0 for backtracking");
            }

            if (isPredictiveCloser(action)) {
                reward += 2.0;
                System.out.println("Reward: +2.0 for predictive move closer to food");
            }

            move();
            double newDistance = getDistanceToFood(snakeX.get(0), snakeY.get(0));
            if (!foodEaten) {
                if (newDistance < oldDistance) {
                    reward += DISTANCE_REWARD_SCALE / newDistance;
                    System.out.println("Reward: +" + String.format("%.2f", DISTANCE_REWARD_SCALE / newDistance) + " for moving closer to food (distance: " + newDistance + ")");
                    consecutiveFoodMoves++;
                    if (consecutiveFoodMoves >= 2) {
                        reward += 1.0;
                        System.out.println("Reward: +1.0 for consecutive move toward food");
                    }
                } else if (newDistance > oldDistance) {
                    reward -= DISTANCE_REWARD_SCALE / newDistance;
                    System.out.println("Penalty: -" + String.format("%.2f", DISTANCE_REWARD_SCALE / newDistance) + " for moving away from food (distance: " + newDistance + ")");
                    consecutiveFoodMoves = 0;
                } else {
                    consecutiveFoodMoves = 0;
                }
            } else {
                reward += 100.0;
                System.out.println("Reward: +100.0 for eating food");
                consecutiveFoodMoves = 0;
            }

            int freeSpace = getFreeSpaceCount();
            if (freeSpace >= 3) {
                reward += 0.5 * freeSpace;
                System.out.println("Reward: +" + String.format("%.2f", 0.5 * freeSpace) + " for free space (" + freeSpace + " cells)");
            } else if (freeSpace <= 1) {
                reward -= 0.5 * (4 - freeSpace);
                System.out.println("Penalty: -" + String.format("%.2f", 0.5 * (4 - freeSpace)) + " for tight space (" + freeSpace + " cells)");
            }

            if (!hasSafePath()) {
                reward -= 0.5;
                System.out.println("Penalty: -0.5 for no safe path to food");
            }

            double[] nextState = getState();
            double collisionReward = checkCollision();
            if (!running) {
                reward = collisionReward;
            }
            trainDQN(state, action, reward, nextState, !running);
            checkFood();
        }
        repaint();
    }

    private class MyKeyAdapter extends KeyAdapter {
        @Override
        public void keyPressed(KeyEvent e) {
            switch (e.getKeyCode()) {
                case KeyEvent.VK_UP:
                    if (direction != 'D') direction = 'U';
                    break;
                case KeyEvent.VK_DOWN:
                    if (direction != 'U') direction = 'D';
                    break;
                case KeyEvent.VK_LEFT:
                    if (direction != 'R') direction = 'L';
                    break;
                case KeyEvent.VK_RIGHT:
                    if (direction != 'L') direction = 'R';
                    break;
                case KeyEvent.VK_SPACE:
                    if (!running) startGame();
                    break;
                case KeyEvent.VK_R:
                    if (!running) startGame();
                    break;
            }
        }
    }

    private class Transition {
        double[] state;
        int action;
        double reward;
        double[] nextState;
        boolean done;

        Transition(double[] state, int action, double reward, double[] nextState, boolean done) {
            this.state = state;
            this.action = action;
            this.reward = reward;
            this.nextState = nextState;
            this.done = done;
        }
    }

    public static void main(String[] args) {
        JFrame frame = new JFrame("Snake Game with DQN");
        SnakeGameWithDQN game = new SnakeGameWithDQN();
        frame.add(game);
        frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        frame.pack();
        frame.setLocationRelativeTo(null);
        frame.setVisible(true);
    }
}